from gm.api import *
import statsmodels.api as sm
import numpy as np
from MSCI_tools import msci_tools


# 斜率越大，支撑相对强弱大，反之则小
def get_rsrs_weight(symbol, now, length=600, window=20):
    last_day = get_previous_trading_date("SZSE", now)
    last_last_day = get_previous_trading_date("SZSE", last_day)
    data = history_n(symbol=symbol, frequency="1d", count=length + window,
                     end_time=last_last_day,
                     fields="open,high,low,close",
                     fill_missing="last", df=True)
    data_high = (data["high"].values)
    data_low = (data["low"].values)

    # 这里是直到前天的RSRS序列值
    rsrs_weights = []

    for len_ in range(len(data) - window + 1):
        high_ = data_high[len_:len_ + window]
        low_ = data_low[len_:len_ + window]

        low_ = sm.add_constant(low_)
        model_ = sm.OLS(high_, low_)
        results_ = model_.fit()
        weight_ = (results_.params[1])

        # rsrs_weights.append(weight_)   #原始weight
        rsrs_weights.append(weight_)  # 这里是加上权重的weight

    last_day_data = history_n(symbol=symbol, frequency="1d", count=window, end_time=last_day,
                              fields="open,high,low,close", fill_missing="last", df=True)

    high_ = last_day_data["high"].values
    low_ = last_day_data["low"].values

    coefficient_determination = (np.corrcoef(high_, low_))[0][1] ** 2

    low_ = sm.add_constant(low_)
    model_ = sm.OLS(high_, low_)
    results_ = model_.fit()
    weight_ = (results_.params[1])

    z_score = (weight_ - np.mean(rsrs_weights)) / np.std(rsrs_weights)

    # RSRS得分
    RSRS_socre = z_score * coefficient_determination

    return RSRS_socre


def get_rsrs_weight_classic(symbol, now, length=200, window=18):
    last_day = get_previous_trading_date('SHSE', now)

    ans = []
    ans_rightdev = []

    data = history_n(symbol, frequency='1d', count=length, end_time=last_day, fields='high,low')

    highs = msci_tools.get_data_value(data, 'high')
    lows = msci_tools.get_data_value(data, 'low')
    for i in range(len(highs))[window:]:
        data_high = highs[i - window + 1:i + 1]
        data_low = lows[i - window + 1:i + 1]
        X = sm.add_constant(data_low)
        model = sm.OLS(data_high, X)
        results = model.fit()
        ans.append(results.params[1])
        # 计算r2
        ans_rightdev.append(results.rsquared)

    # 计算标准化的RSRS指标
    # 计算均值序列
    section = ans[-length:]
    # 计算均值序列
    mu = np.mean(section)
    # 计算标准化右偏RSRS标准分
    sigma = np.std(section)
    zscore = (section[-1] - mu) / sigma
    # 计算右偏RSRS标准分
    zscore_rightdev = zscore * ans[-1] * ans_rightdev[-1]

    return (zscore_rightdev)
